/* Copyright (C) 2005, 2006, 2007, 2008, 2009, 2010 Mario Orsi
   This file is part of Brahms.
   Brahms is free software: you can redistribute it and/or modify it under the terms of the GNU General Public License as published by
   the Free Software Foundation, either version 3 of the License, or (at your option) any later version.
   Brahms is distributed in the hope that it will be useful, but WITHOUT ANY WARRANTY; without even the implied warranty of
   MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU General Public License for more details.
   You should have received a copy of the GNU General Public License along with Brahms.  If not, see <http://www.gnu.org/licenses/>. */

/* rotation.c -- functions used to operate on vectors and matrices.
   Note: a matrix is treated as a linear array, e.g.:
   M = m[0] m[3] m[6]
   m[1] m[4] m[7]
   m[2] m[5] m[8]
*/

#include "dataStructs.h"

void BuildRotMatrixFromQuat( RMat *rMat, const Quat *q, const int transpose ) // [Rapaport,pp. 207-208]
{
  real p[ 10 ], tq[ 4 ], s;
  int k, k1, k2;

  tq[ 0 ] = q->u1;
  tq[ 1 ] = q->u2;
  tq[ 2 ] = q->u3;
  tq[ 3 ] = q->u4;
  for ( k = 0, k2 = 0; k2 < 4; k2++ ) {
    for ( k1 = k2; k1 < 4; k1++, k++ )
      p[ k ] = 2. * tq[ k1 ] * tq[ k2 ];
  }
  rMat->u[0] = p[0] + p[9] - 1.;
  rMat->u[4] = p[4] + p[9] - 1.;
  rMat->u[8] = p[7] + p[9] - 1.;
  s = transpose ? 1. : -1.;
  rMat->u[1] = p[1] + s * p[8];
  rMat->u[3] = p[1] - s * p[8];
  rMat->u[2] = p[2] - s * p[6];
  rMat->u[6] = p[2] + s * p[6];
  rMat->u[5] = p[5] + s * p[3];
  rMat->u[7] = p[5] - s * p[3];
}

void BuildRxMatrix( RMat *rMat, const real angle ) 
{
  real angleSq = Sqr( angle );
  real cosAngle = ( 1. - angleSq / 4. ) / ( 1. + angleSq / 4. );
  real sinAngle = angle / ( 1. + angleSq / 4. );

  rMat->u[ 0 ] = 1.;
  rMat->u[ 3 ] = 0.;
  rMat->u[ 6 ] = 0.;
  rMat->u[ 1 ] = 0.;
  rMat->u[ 4 ] = cosAngle;
  rMat->u[ 7 ] = sinAngle;
  rMat->u[ 2 ] = 0.;
  rMat->u[ 5 ] = - sinAngle;
  rMat->u[ 8 ] = cosAngle;
}

void BuildRyMatrix( RMat *rMat, const real angle ) 
{
  real angleSq = Sqr( angle );
  real cosAngle = ( 1. - angleSq / 4. ) / ( 1. + angleSq / 4. );
  real sinAngle = angle / ( 1. + angleSq / 4. );

  rMat->u[ 0 ] = cosAngle;
  rMat->u[ 3 ] = 0.;
  rMat->u[ 6 ] = - sinAngle;
  rMat->u[ 1 ] = 0.;
  rMat->u[ 4 ] = 1.;
  rMat->u[ 7 ] = 0.;
  rMat->u[ 2 ] = sinAngle;
  rMat->u[ 5 ] = 0.;
  rMat->u[ 8 ] = cosAngle;
}

void BuildRzMatrix( RMat *rMat, const real angle ) 
{
  real angleSq = Sqr( angle );
  real cosAngle = ( 1. - angleSq / 4. ) / ( 1. + angleSq / 4. );
  real sinAngle = angle / ( 1. + angleSq / 4. );

  rMat->u[ 0 ] = cosAngle;
  rMat->u[ 3 ] = sinAngle;
  rMat->u[ 6 ] = 0.;
  rMat->u[ 1 ] = - sinAngle;
  rMat->u[ 4 ] = cosAngle;
  rMat->u[ 7 ] = 0.;
  rMat->u[ 2 ] = 0.;                
  rMat->u[ 5 ] = 0.;
  rMat->u[ 8 ] = 1.;
}

void UpdateQMatrix( RMat *qMat, const RMat *rMat ) 
{
  RMat tMat, qMatTrans; /* temporary matrices */
  tMat.u[0] = qMat->u[0]; /* storing Q^T in temporary matrix */
  tMat.u[1] = qMat->u[3];
  tMat.u[2] = qMat->u[6];
  tMat.u[3] = qMat->u[1];
  tMat.u[4] = qMat->u[4];
  tMat.u[5] = qMat->u[7];
  tMat.u[6] = qMat->u[2];
  tMat.u[7] = qMat->u[5];
  tMat.u[8] = qMat->u[8];

  qMatTrans.u[0] = tMat.u[0] * rMat->u[0] + tMat.u[3] * rMat->u[3] + tMat.u[6] * rMat->u[6]; /* Q^T := Q^T * R^T  */
  qMatTrans.u[1] = tMat.u[1] * rMat->u[0] + tMat.u[4] * rMat->u[3] + tMat.u[7] * rMat->u[6];
  qMatTrans.u[2] = tMat.u[2] * rMat->u[0] + tMat.u[5] * rMat->u[3] + tMat.u[8] * rMat->u[6];
  qMatTrans.u[3] = tMat.u[0] * rMat->u[1] + tMat.u[3] * rMat->u[4] + tMat.u[6] * rMat->u[7];
  qMatTrans.u[4] = tMat.u[1] * rMat->u[1] + tMat.u[4] * rMat->u[4] + tMat.u[7] * rMat->u[7];
  qMatTrans.u[5] = tMat.u[2] * rMat->u[1] + tMat.u[5] * rMat->u[4] + tMat.u[8] * rMat->u[7];
  qMatTrans.u[6] = tMat.u[0] * rMat->u[2] + tMat.u[3] * rMat->u[5] + tMat.u[6] * rMat->u[8];
  qMatTrans.u[7] = tMat.u[1] * rMat->u[2] + tMat.u[4] * rMat->u[5] + tMat.u[7] * rMat->u[8];
  qMatTrans.u[8] = tMat.u[2] * rMat->u[2] + tMat.u[5] * rMat->u[5] + tMat.u[8] * rMat->u[8];

  qMat->u[0] = qMatTrans.u[0]; /* Q := (Q^T)^T */
  qMat->u[1] = qMatTrans.u[3];
  qMat->u[2] = qMatTrans.u[6];
  qMat->u[3] = qMatTrans.u[1];
  qMat->u[4] = qMatTrans.u[4];
  qMat->u[5] = qMatTrans.u[7];
  qMat->u[6] = qMatTrans.u[2];
  qMat->u[7] = qMatTrans.u[5];
  qMat->u[8] = qMatTrans.u[8];
}

void GetRotMatFromOrientVector( RMat *R, const VecR *orient ) 
{
  R->u[ 0 ] = 0.; 
  R->u[ 1 ] = 0.;
  R->u[ 2 ] = orient->x;
  R->u[ 3 ] = 0.;
  R->u[ 4 ] = 0.;
  R->u[ 5 ] = orient->y;
  R->u[ 6 ] = 0.;
  R->u[ 7 ] = 0.;
  R->u[ 8 ] = orient->z;
}

void GetOrientPrincAxisX_FromRotMat( VecR *orient, const RMat R ) 
{
  VecR refOrient;
  VSet( refOrient, 1., 0., 0. ); // lab frame x-axis
  MVMulT( *orient, R.u, refOrient ); // body (aka principal) frame x-axis
}

void GetOrientPrincAxisY_FromRotMat( VecR *orient, const RMat R ) 
{
  VecR refOrient;
  VSet( refOrient, 0., 1., 0. ); // lab frame y-axis
  MVMulT( *orient, R.u, refOrient ); // body (aka principal) frame y-axis
}

void GetOrientPrincAxisZ_FromRotMat( VecR *orient, const RMat R ) 
{
  VecR refOrient;  
  VSet( refOrient, 0., 0., 1. ); // lab frame z-axis
  MVMulT( *orient, R.u, refOrient ); // body (aka principal) frame z-axis
}

void EulerToRotMat( RMat *R, const real phi, const real theta, const real psi ) 
{ // ref: Allen & Tildesley, section 3.3.1, page 86
  real cphi, ctheta, cpsi, sphi, stheta, spsi;
  cphi = cos( phi );
  sphi = sin( phi );
  ctheta = cos( theta );
  stheta = sin( theta );
  cpsi = cos( psi );
  spsi = sin( psi );
  R->u[ 0 ] =  cphi * cpsi - sphi * ctheta * spsi;
  R->u[ 1 ] = -cphi * spsi - sphi * ctheta * cpsi;
  R->u[ 2 ] =  sphi * stheta;
  R->u[ 3 ] =  sphi * cpsi + cphi * ctheta * spsi;
  R->u[ 4 ] = -sphi * spsi + cphi * ctheta * cpsi;
  R->u[ 5 ] = -cphi * stheta;
  R->u[ 6 ] =  stheta * spsi;
  R->u[ 7 ] =  stheta * cpsi;
  R->u[ 8 ] =  ctheta;
}

void RotMatToQuat( Quat *quat, const RMat *R ) 
{ // input: R - output: quaternion (NOT SURE THIS IS CORRECT)
  real w, x, y, z; // quat components
  real w4; // 4 * w
  w = 0.5 * sqrt( 1. + R->u[ 0 ] + R->u[ 4 ] + R->u[ 8 ] ); 
  w4 =  4. * w;
  x = ( R->u[ 5 ] - R->u[ 7 ] ) / w4;
  y = ( R->u[ 6 ] - R->u[ 2 ] ) / w4;
  z = ( R->u[ 1 ] - R->u[ 3 ] ) / w4;
  QSet (*quat, -x, -y, -z, w );
}

void BuildQuatFromRotMat( Quat *quat, const RMat *R )
{
  real t, s, q[4];
  real ad1, ad2, ad3;
    
  t = R->u[ 0 ] + R->u[ 4 ] + R->u[ 8 ] + 1.0;
  if( t > 0.0 ){
    
    s = 0.5 / sqrt( t );
    q[0] = 0.25 / s;
    q[1] = (R->u[ 7 ] - R->u[ 5 ]) * s;
    q[2] = (R->u[ 2 ] - R->u[ 6 ]) * s;
    q[3] = (R->u[ 3 ] - R->u[ 1 ]) * s;
  }
  else{
    
    ad1 = fabs( R->u[ 0 ] );
    ad2 = fabs( R->u[ 4 ] );
    ad3 = fabs( R->u[ 8 ] );
    
    if( ad1 >= ad2 && ad1 >= ad3 ){
      
      s = 2.0 * sqrt( 1.0 + R->u[ 0 ] - R->u[ 4 ] - R->u[ 8 ] );
      q[0] = (R->u[ 7 ] + R->u[ 5 ]) / s;
      q[1] = 0.5 / s;
      q[2] = (R->u[ 3 ] + R->u[ 1 ]) / s;
      q[3] = (R->u[ 6 ] + R->u[ 2 ]) / s;
    }
    else if( ad2 >= ad1 && ad2 >= ad3 ){
      
      s = sqrt( 1.0 + R->u[ 4 ] - R->u[ 0 ] - R->u[ 8 ] ) * 2.0;
      q[0] = (R->u[ 6 ] + R->u[ 2 ]) / s;
      q[1] = (R->u[ 3 ] + R->u[ 1 ]) / s;
      q[2] = 0.5 / s;
      q[3] = (R->u[ 7 ] + R->u[ 5 ]) / s;
    }
    else{
  
      s = sqrt( 1.0 + R->u[ 8 ] - R->u[ 0 ] - R->u[ 4 ] ) * 2.0;
      q[0] = (R->u[ 3 ] + R->u[ 1 ]) / s;
      q[1] = (R->u[ 6 ] + R->u[ 2 ]) / s;
      q[2] = (R->u[ 7 ] + R->u[ 5 ]) / s;
      q[3] = 0.5 / s;
    }
  }
  QSet( *quat, q[1], q[2], q[3], q[0] );
}


void QuatToEuler ( real *phi, real *theta, real *psi, const Quat *qe )
{
  *phi = atan( 2. * ( qe->u4 * qe->u1 + qe->u2 * qe->u3 ) / ( 1. - 2. * ( Sqr( qe->u1 ) + Sqr( qe->u2 ) ) ) ); // phi
  *theta = asin( 2. * ( qe->u4 * qe->u2 - qe->u3 * qe->u1 ) ); // theta
  *psi = atan( 2. * ( qe->u4 * qe->u3 + qe->u1 * qe->u2 ) / ( 1. - 2. * ( Sqr( qe->u2 ) + Sqr( qe->u3 ) ) ) ); // psi
}

void EulerToQuat (Quat *qe, const real *eAng) // Rapaport, The Art of MD (2004), p215.
{
  real a1, a2, a3;
  a1 = 0.5 * eAng[1];
  a2 = 0.5 * (eAng[0] - eAng[2]);
  a3 = 0.5 * (eAng[0] + eAng[2]);
  QSet (*qe, sin (a1) * cos (a2), sin (a1) * sin (a2), cos (a1) * sin (a3), cos (a1) * cos (a3));
}
